{
 "cells": [
  {
   "cell_type": "code",
   "id": "initial_id",
   "metadata": {
    "collapsed": true,
    "ExecuteTime": {
     "end_time": "2025-09-17T15:04:19.355901Z",
     "start_time": "2025-09-17T15:03:58.590794Z"
    }
   },
   "source": [
    "from llama_index.finetuning import EmbeddingQAFinetuneDataset\n",
    "from datasets import load_dataset\n",
    "\n",
    "\n",
    "def load_hf_dataset(dataset_name):\n",
    "    hf_dataset_name = f\"sepz/{dataset_name}_ft\"\n",
    "    corpus = load_dataset(hf_dataset_name, \"data_records\", split=\"train\")\n",
    "\n",
    "    queries_train = load_dataset(hf_dataset_name, \"qs\", split=\"train\")\n",
    "    queries_validation = load_dataset(hf_dataset_name, \"qs\", split=\"dev\")\n",
    "    queries_test = load_dataset(hf_dataset_name, \"qs\", split=\"test\")\n",
    "\n",
    "    qrels_train = load_dataset(hf_dataset_name, \"qs_rel\", split=\"train\")\n",
    "    qrels_validation = load_dataset(hf_dataset_name, \"qs_rel\", split=\"dev\")\n",
    "    qrels_test = load_dataset(hf_dataset_name, \"qs_rel\", split=\"test\")\n",
    "\n",
    "    corpus = {\n",
    "        str(corpus[i][\"record_id\"]): corpus[i][\"text\"]\n",
    "        for i in range(len(corpus))\n",
    "    }\n",
    "\n",
    "    queries_train = {\n",
    "        str(queries_train[i][\"q_id\"]): queries_train[i][\"input\"]\n",
    "        for i in range(len(queries_train))\n",
    "    }\n",
    "    queries_validation = {\n",
    "        str(r[\"q_id\"]): r[\"input\"] for r in queries_validation\n",
    "    }\n",
    "    queries_test = {str(r[\"q_id\"]): r[\"input\"] for r in queries_test}\n",
    "\n",
    "    qrels_train = (\n",
    "        qrels_train.to_pandas()\n",
    "        .groupby(\"q_id\")[\"record_id\"]\n",
    "        .apply(list)\n",
    "        .to_dict()\n",
    "    )\n",
    "    qrels_validation = (\n",
    "        qrels_validation.to_pandas()\n",
    "        .groupby(\"q_id\")[\"record_id\"]\n",
    "        .apply(list)\n",
    "        .to_dict()\n",
    "    )\n",
    "    qrels_test = (\n",
    "        qrels_test.to_pandas()\n",
    "        .groupby(\"q_id\")[\"record_id\"]\n",
    "        .apply(list)\n",
    "        .to_dict()\n",
    "    )\n",
    "    # convert to strings\n",
    "    qrels_train = {str(k): [str(i) for i in v] for k, v in qrels_train.items()}\n",
    "    qrels_validation = {\n",
    "        str(k): [str(i) for i in v] for k, v in qrels_validation.items()\n",
    "    }\n",
    "    qrels_test = {str(k): [str(i) for i in v] for k, v in qrels_test.items()}\n",
    "\n",
    "    # Load the dataset\n",
    "    train_dataset = EmbeddingQAFinetuneDataset(\n",
    "        corpus=corpus, queries=queries_train, relevant_docs=qrels_train\n",
    "    )\n",
    "    validation_dataset = EmbeddingQAFinetuneDataset(\n",
    "        corpus=corpus,\n",
    "        queries=queries_validation,\n",
    "        relevant_docs=qrels_validation,\n",
    "    )\n",
    "    test_dataset = EmbeddingQAFinetuneDataset(\n",
    "        corpus=corpus, queries=queries_test, relevant_docs=qrels_test\n",
    "    )\n",
    "\n",
    "    return train_dataset, validation_dataset, test_dataset"
   ],
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:datasets:PyTorch version 2.5.1+cu121 available.\n",
      "PyTorch version 2.5.1+cu121 available.\n"
     ]
    }
   ],
   "execution_count": 1
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-09-17T15:06:33.805044Z",
     "start_time": "2025-09-17T15:04:21.847507Z"
    }
   },
   "cell_type": "code",
   "source": [
    "from llama_index.core.embeddings import resolve_embed_model\n",
    "\n",
    "train_dataset, val_dataset, test_dataset = load_hf_dataset(\"scifact\")\n",
    "base_embed_model = resolve_embed_model(\"local:BAAI/bge-small-en-v1.5\")"
   ],
   "id": "2f17c3f97f931f5a",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "README.md:   0%|          | 0.00/814 [00:00<?, ?B/s]"
      ],
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "19d0c9f9456347dfabb92aeccf72d7e4"
      }
     },
     "metadata": {},
     "output_type": "display_data",
     "jetTransient": {
      "display_id": null
     }
    },
    {
     "data": {
      "text/plain": [
       "data.parquet:   0%|          | 0.00/4.57M [00:00<?, ?B/s]"
      ],
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "d681d36ae95e4dc3bf80ba47c0816db2"
      }
     },
     "metadata": {},
     "output_type": "display_data",
     "jetTransient": {
      "display_id": null
     }
    },
    {
     "data": {
      "text/plain": [
       "Generating train split:   0%|          | 0/5183 [00:00<?, ? examples/s]"
      ],
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "bb268cc7e95c40c08588d498f48f680d"
      }
     },
     "metadata": {},
     "output_type": "display_data",
     "jetTransient": {
      "display_id": null
     }
    },
    {
     "data": {
      "text/plain": [
       "Generating dev split:   0%|          | 0/5183 [00:00<?, ? examples/s]"
      ],
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "e404e8de87ab49a2bc57a4b5c13bcd78"
      }
     },
     "metadata": {},
     "output_type": "display_data",
     "jetTransient": {
      "display_id": null
     }
    },
    {
     "data": {
      "text/plain": [
       "Generating test split:   0%|          | 0/5183 [00:00<?, ? examples/s]"
      ],
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "f06f76c9fe924748aaba1fca86a07bd7"
      }
     },
     "metadata": {},
     "output_type": "display_data",
     "jetTransient": {
      "display_id": null
     }
    },
    {
     "data": {
      "text/plain": [
       "qs.parquet:   0%|          | 0.00/48.0k [00:00<?, ?B/s]"
      ],
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "63675251a94c410896cba05daf75e7fd"
      }
     },
     "metadata": {},
     "output_type": "display_data",
     "jetTransient": {
      "display_id": null
     }
    },
    {
     "data": {
      "text/plain": [
       "qs.parquet:   0%|          | 0.00/10.9k [00:00<?, ?B/s]"
      ],
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "c0597a5b1e3747c79365f66dcad84e3e"
      }
     },
     "metadata": {},
     "output_type": "display_data",
     "jetTransient": {
      "display_id": null
     }
    },
    {
     "data": {
      "text/plain": [
       "qs.parquet:   0%|          | 0.00/17.7k [00:00<?, ?B/s]"
      ],
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "ab5f94a4bad74d9a8611944b972a46b6"
      }
     },
     "metadata": {},
     "output_type": "display_data",
     "jetTransient": {
      "display_id": null
     }
    },
    {
     "data": {
      "text/plain": [
       "Generating train split:   0%|          | 0/778 [00:00<?, ? examples/s]"
      ],
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "8835b74e5e5c478cad8026a98da2894c"
      }
     },
     "metadata": {},
     "output_type": "display_data",
     "jetTransient": {
      "display_id": null
     }
    },
    {
     "data": {
      "text/plain": [
       "Generating dev split:   0%|          | 0/110 [00:00<?, ? examples/s]"
      ],
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "ca5b2708988946b5a1379765878a6859"
      }
     },
     "metadata": {},
     "output_type": "display_data",
     "jetTransient": {
      "display_id": null
     }
    },
    {
     "data": {
      "text/plain": [
       "Generating test split:   0%|          | 0/221 [00:00<?, ? examples/s]"
      ],
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "f6509dcba96448b6858d9ca86533e336"
      }
     },
     "metadata": {},
     "output_type": "display_data",
     "jetTransient": {
      "display_id": null
     }
    },
    {
     "data": {
      "text/plain": [
       "qs_rel.parquet:   0%|          | 0.00/9.96k [00:00<?, ?B/s]"
      ],
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "c7244b7773f14de69d07643ea86e4b5a"
      }
     },
     "metadata": {},
     "output_type": "display_data",
     "jetTransient": {
      "display_id": null
     }
    },
    {
     "data": {
      "text/plain": [
       "qs_rel.parquet:   0%|          | 0.00/3.70k [00:00<?, ?B/s]"
      ],
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "7a439bdc30c34701bd1c2908d9249da9"
      }
     },
     "metadata": {},
     "output_type": "display_data",
     "jetTransient": {
      "display_id": null
     }
    },
    {
     "data": {
      "text/plain": [
       "qs_rel.parquet:   0%|          | 0.00/4.87k [00:00<?, ?B/s]"
      ],
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "fb61b60f79854143a78e9135a1153e3c"
      }
     },
     "metadata": {},
     "output_type": "display_data",
     "jetTransient": {
      "display_id": null
     }
    },
    {
     "data": {
      "text/plain": [
       "Generating train split:   0%|          | 0/868 [00:00<?, ? examples/s]"
      ],
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "696ce250406841e4866fffa3e609ba76"
      }
     },
     "metadata": {},
     "output_type": "display_data",
     "jetTransient": {
      "display_id": null
     }
    },
    {
     "data": {
      "text/plain": [
       "Generating dev split:   0%|          | 0/129 [00:00<?, ? examples/s]"
      ],
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "f8288336cc194074bb253253316ed0e5"
      }
     },
     "metadata": {},
     "output_type": "display_data",
     "jetTransient": {
      "display_id": null
     }
    },
    {
     "data": {
      "text/plain": [
       "Generating test split:   0%|          | 0/261 [00:00<?, ? examples/s]"
      ],
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "1a3ed1f836674a96aebce937402043b3"
      }
     },
     "metadata": {},
     "output_type": "display_data",
     "jetTransient": {
      "display_id": null
     }
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:sentence_transformers.SentenceTransformer:Load pretrained SentenceTransformer: BAAI/bge-small-en-v1.5\n",
      "Load pretrained SentenceTransformer: BAAI/bge-small-en-v1.5\n",
      "INFO:sentence_transformers.SentenceTransformer:1 prompt is loaded, with the key: query\n",
      "1 prompt is loaded, with the key: query\n"
     ]
    }
   ],
   "execution_count": 2
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-09-17T15:06:33.932566Z",
     "start_time": "2025-09-17T15:06:33.906570Z"
    }
   },
   "cell_type": "code",
   "source": "print(val_dataset.queries[\"2\"])",
   "id": "34935785c28e4bee",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Depletion of nitric oxide is responsible for vasospasm.\n"
     ]
    }
   ],
   "execution_count": 3
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-09-17T15:06:34.059684Z",
     "start_time": "2025-09-17T15:06:34.038688Z"
    }
   },
   "cell_type": "code",
   "source": [
    "# This code would generate your own dataset against your own custom data\n",
    "from llama_index.finetuning import generate_qa_embedding_pairs\n",
    "from llama_index.core import SimpleDirectoryReader\n",
    "from llama_index.core.node_parser import SentenceSplitter\n",
    "from llama_index.core.evaluation import EmbeddingQAFinetuneDataset\n",
    "\n",
    "\n",
    "def load_corpus(files, verbose=False):\n",
    "    if verbose:\n",
    "        print(f\"Loading files {files}\")\n",
    "\n",
    "    reader = SimpleDirectoryReader(input_files=files)\n",
    "    docs = reader.load_data()\n",
    "    if verbose:\n",
    "        print(f\"Loaded {len(docs)} docs\")\n",
    "\n",
    "    parser = SentenceSplitter()\n",
    "    nodes = parser.get_nodes_from_documents(docs, show_progress=verbose)\n",
    "\n",
    "    if verbose:\n",
    "        print(f\"Parsed {len(nodes)} nodes\")\n",
    "\n",
    "    return nodes\n",
    "\n",
    "\n",
    "# Load data\n",
    "# train_nodes = load_corpus([\"file1.pdf\", ...], verbose=True)\n",
    "# val_nodes = load_corpus([\"file2.pdf\", ...], verbose=True)\n",
    "\n",
    "# Generate pairs\n",
    "# train_dataset = generate_qa_embedding_pairs(train_nodes)\n",
    "# val_dataset = generate_qa_embedding_pairs(val_nodes)\n",
    "\n",
    "# [Optional] Save to disk\n",
    "# train_dataset.save_json(\"train_dataset.json\")\n",
    "# val_dataset.save_json(\"val_dataset.json\")\n",
    "\n",
    "# [Optional] Load\n",
    "# train_dataset = EmbeddingQAFinetuneDataset.from_json(\"train_dataset.json\")\n",
    "# val_dataset = EmbeddingQAFinetuneDataset.from_json(\"val_dataset.json\")"
   ],
   "id": "3f631b2b7476dbee",
   "outputs": [],
   "execution_count": 4
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-09-17T15:06:58.824830Z",
     "start_time": "2025-09-17T15:06:58.797832Z"
    }
   },
   "cell_type": "code",
   "source": [
    "from typing import Optional, Dict\n",
    "\n",
    "import torch\n",
    "import numpy as np\n",
    "from tqdm import tqdm\n",
    "from llama_index.core.schema import TextNode\n",
    "from llama_index.core.base.embeddings.base import BaseEmbedding\n",
    "from llama_index.core.base.base_retriever import BaseRetriever\n",
    "from llama_index.core import VectorStoreIndex\n",
    "\n",
    "\n",
    "def build_retriever(\n",
    "    corpus: Dict[str, str],\n",
    "    embed_model: BaseEmbedding | str,\n",
    "    corpus_embeddings: Optional[torch.Tensor] = None,\n",
    "    k: int = 10,\n",
    ") -> BaseRetriever:\n",
    "    nodes = []\n",
    "    for i, (id_, text) in enumerate(corpus.items()):\n",
    "        if corpus_embeddings is not None:\n",
    "            nodes.append(\n",
    "                TextNode(\n",
    "                    id_=id_, text=text, embedding=corpus_embeddings[i].tolist()\n",
    "                )\n",
    "            )\n",
    "        else:\n",
    "            nodes.append(TextNode(id_=id_, text=text))\n",
    "\n",
    "    index = VectorStoreIndex(\n",
    "        nodes=nodes,\n",
    "        embeddings=corpus_embeddings,\n",
    "        embed_model=embed_model,\n",
    "        show_progress=True,\n",
    "    )\n",
    "    return index.as_retriever(similarity_top_k=k)\n",
    "\n",
    "\n",
    "def ndcg_at_k(\n",
    "    dataset: EmbeddingQAFinetuneDataset, retriever: BaseRetriever, k: int = 10\n",
    "):\n",
    "    queries = dataset.queries\n",
    "    relevant_docs = dataset.relevant_docs\n",
    "    ndcg_scores = []\n",
    "    for query_id, query in tqdm(queries.items()):\n",
    "        retrieved_nodes = retriever.retrieve(query)\n",
    "        retrieved_ids = [node.node.node_id for node in retrieved_nodes]\n",
    "        expected_ids = relevant_docs[query_id]\n",
    "\n",
    "        # Calculate NDCG\n",
    "        ideal_dcg = np.sum(\n",
    "            [1 / np.log2(i + 2) for i in range(min(k, len(expected_ids)))]\n",
    "        )\n",
    "        rel_scores = np.zeros(k)\n",
    "        for j in range(min(k, len(retrieved_ids))):\n",
    "            if retrieved_ids[j] in expected_ids:\n",
    "                rel_scores[j] = 1\n",
    "        dcg = np.sum(\n",
    "            [rel_scores[i] / np.log2(i + 2) for i in range(len(rel_scores))]\n",
    "        )\n",
    "        ndcg = dcg / ideal_dcg if ideal_dcg > 0 else 0\n",
    "\n",
    "        ndcg_scores.append(ndcg)\n",
    "\n",
    "    mean_ndcg = np.mean(ndcg_scores)\n",
    "    return mean_ndcg"
   ],
   "id": "17f66ddb67a37742",
   "outputs": [],
   "execution_count": 6
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-09-17T15:23:20.470406Z",
     "start_time": "2025-09-17T15:18:37.245470Z"
    }
   },
   "cell_type": "code",
   "source": [
    "%%capture\n",
    "from llama_index.experimental import Nudge\n",
    "\n",
    "k = 10\n",
    "\n",
    "nudge = Nudge(\n",
    "    train_dataset=train_dataset,\n",
    "    val_dataset=val_dataset,\n",
    "    embed_model=base_embed_model,\n",
    "    use_nudge_n=True,\n",
    ")\n",
    "nudge.finetune()\n",
    "nudge_corpus_embeddings = nudge.get_finetuned_corpus_embeddings()\n",
    "nudge_retriever = build_retriever(\n",
    "    train_dataset.corpus, base_embed_model, nudge_corpus_embeddings, k=k\n",
    ")\n",
    "nudge_ndcg_test = ndcg_at_k(test_dataset, nudge_retriever, k)"
   ],
   "id": "b746bd61799a8eb8",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:llama_index.experimental.nudge.base:Use pytorch device: cuda\n",
      "Use pytorch device: cuda\n"
     ]
    }
   ],
   "execution_count": 9
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": "",
   "id": "cec835f1b362a9e"
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
